
#include "monte_carlo/engine.hpp"

#include "information/affected_range.hpp"
#include "information/information_none.hpp"
#include "information/information_exact.hpp"
#include "information/information_at_most.hpp"

#include "known_thresholds.hpp"
#include "matlab.hpp"

#include "statistic/bayes_cusum.hpp"
#include "statistic/generalized_cusum.hpp"
#include "statistic/oracle_cusum.hpp"
#include "statistic/ordered_sum_cusum.hpp"
#include "statistic/parallel_max_cusum.hpp"
#include "statistic/parallel_mix_cusum.hpp"

#include "window/window_all_cusum.hpp"
#include "window/window_fixed.hpp"
#include "window/window_none.hpp"
#include "window/window_rlr.hpp"

#include "unit_test/test_distribution.hpp"
#include "unit_test/test_information.hpp"

#include <array>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <functional>
#include <iostream>
#include <random>
#include <string>
#include <vector>

using namespace ropufu;

template <typename t_statistic_type>
void find_arl(std::uint_fast64_t n, t_statistic_type& proc, double shift_in_mean, double arl);

template <typename t_statistic_type>
void sequence_thresholds(std::uint_fast64_t n, t_statistic_type& proc, std::size_t count_affected,
    double shift_in_mean, double arl_from, double arl_to, std::size_t count_thresholds);

template <typename t_statistic_type, std::size_t t_count_possibly_affected>
struct fork
{
    typedef t_statistic_type statistic_type;
    typedef typename t_statistic_type::information_type information_type;

    static void execute(std::uint_fast64_t n, t_statistic_type& proc, double shift_in_mean, double arl_from, double arl_to, std::size_t count_thresholds)
    {
        sequence_thresholds(n, proc, 0, shift_in_mean, arl_from, arl_to, count_thresholds); // ARL.
        for (std::size_t k = information::affected_range<information_type>::value_from; k <= information::affected_range<information_type>::value_to; k++)
        {
            sequence_thresholds(n, proc, k, shift_in_mean, arl_from, arl_to, count_thresholds);
        }
    }
};

template <std::size_t t_count_streams, std::size_t t_count_possibly_affected>
struct oracle
{
    typedef statistic::oracle_cusum<t_count_streams, t_count_possibly_affected> oracle_type;

    static void execute(std::uint_fast64_t n, double shift_in_mean, double arl_from, double arl_to, std::size_t count_thresholds)
    {
        std::array<bool, t_count_streams> affected_subset_mask;
        affected_subset_mask.fill(false);
        for (std::size_t k = 0; k < t_count_possibly_affected; k++) affected_subset_mask[k] = true;
        oracle_type proc = oracle_type(affected_subset_mask);

        //find_arl(oracle, shift_in_mean, 10'000UL, arl_from);
        //find_arl(oracle, shift_in_mean, 1'000UL, arl_to);

        sequence_thresholds(n, proc, 0, shift_in_mean, arl_from, arl_to, count_thresholds);
        sequence_thresholds(n, proc, t_count_possibly_affected, shift_in_mean, arl_from, arl_to, count_thresholds);

        oracle<t_count_streams, t_count_possibly_affected - 1>::execute(n, shift_in_mean, arl_from, arl_to, count_thresholds);
    }
};

template <std::size_t t_count_streams>
struct oracle<t_count_streams, 0>
{
    static void execute(std::uint_fast64_t n, double shift_in_mean, double arl_from, double arl_to, std::size_t count_thresholds)
    {
    }
};

template <std::size_t t_count_possibly_affected>
void main_not();

int main() //(std::int32_t argc, char* argv[], char* envp[])
{
    //main_not<5>();
    main_not<4>();
    main_not<3>();
    main_not<2>();
    //main_not<1>();
    return 0;
}

template <std::size_t t_count_possibly_affected>
void main_not()
{
    const std::size_t count_streams = 5;
    const std::size_t count_possibly_affected = t_count_possibly_affected;

    const bool do_find_arl = false;
    const bool do_oracle = false;

    double shift_in_mean = 1.0;
    double arl_from = 1e2;
    double arl_to = 1e5;

    const std::size_t count_thresholds = aftermath::npow<2, 10>::value + 1;
    const std::uint_fast64_t n = 50'000;

    //ropufu::unit_test::test_information<count_streams, count_possibly_affected>();
    //return 0;
    //ropufu::unit_test::test_distribution_normal();
    //return 0;

    typedef information_none<count_streams> info_none_t;
    typedef information_at_most<count_streams, count_possibly_affected> info_at_most_t;
    typedef information_exact<count_streams, count_possibly_affected>   info_exact_t;

    typedef info_exact_t info_type;

    typedef statistic::ordered_sum_cusum<count_streams, info_type> mei_t;
    typedef statistic::parallel_max_cusum<count_streams, info_type> max_cusum_t;
    typedef statistic::parallel_mix_cusum<count_streams, info_type> mix_cusum_t;
    typedef statistic::bayes_cusum<count_streams, info_type> bayes_cusum_t;
    //typedef statistic::generalized_cusum<count_streams, info_type> glr_cusum_t;

    typedef mix_cusum_t proc_type;

    info_type info;
    proc_type proc = proc_type(info);
    //glr_cusum_t proc = glr_cusum_t(info, win_cusum_t());

    if (do_find_arl)
    {
        find_arl(n, proc, shift_in_mean, arl_from);
        find_arl(n / 50, proc, shift_in_mean, arl_to);
    }
    else
    {
        if (do_oracle) oracle<count_streams, count_streams>::execute(n, shift_in_mean, arl_from, arl_to, count_thresholds);
        else fork<proc_type, count_possibly_affected>::execute(n, proc, shift_in_mean, arl_from, arl_to, count_thresholds);
    }
}

template <typename t_statistic_type>
void find_arl(std::uint_fast64_t n, t_statistic_type& proc, double shift_in_mean, double arl)
{
    std::uint_fast32_t seed_override = static_cast<std::uint_fast32_t>(std::chrono::high_resolution_clock::now().time_since_epoch().count());
    seed_override = 0;

    bool is_found = false;

    monte_carlo::engine<t_statistic_type, aftermath::probability::dist_normal> mc_engine(0.0, shift_in_mean, shift_in_mean, seed_override);

    double a = 0.0;
    double step = 1.0;
    std::vector<double> threshold(1);
    std::vector<double> expected_time(1);
    std::vector<double> expected_aux(1);
    while (!is_found)
    {
        threshold[0] = a;
        mc_engine.execute(n, proc, 0, threshold, expected_time, expected_aux, false, false, false); // ARL
        double x = expected_time[0];
        double delta = x - arl;
        if (delta < 0) a += step;
        else
        {
            a -= step;
            step /= 2;
            a += step;
        }
        if (step < 0.01) is_found = true;
    }
}

template <typename t_statistic_type>
void sequence_thresholds(std::uint_fast64_t n, t_statistic_type& proc, std::size_t count_affected,
    double shift_in_mean, double arl_from, double arl_to, std::size_t count_thresholds)
{
    std::uint_fast32_t seed_override = static_cast<std::uint_fast32_t>(std::chrono::high_resolution_clock::now().time_since_epoch().count());
    seed_override = 0;
    const bool is_quiet = false;
    const bool do_dump = true;
    const bool do_stop_on_reset = false;

    bool is_known = true;

    double threshold_from = 0.0;
    double threshold_to = 0.0;
    double dummy;
    is_known &= known_thresholds::for_theta_and_arl(proc, shift_in_mean, arl_from, std::ref(threshold_from), std::ref(dummy));
    is_known &= known_thresholds::for_theta_and_arl(proc, shift_in_mean, arl_to, std::ref(dummy), std::ref(threshold_to));

    if (!is_known) std::cout << "Threshold unknown . . ." << std::endl;
    else
    {
        matlab::tic();
        //std::vector<double> threshold = matlab::expspace(threshold_from, threshold_to, count_thresholds);
        
        // ~~ Fix low end of thresholds ~~
        threshold_to = (0.6) * threshold_from + (0.4) * threshold_to;
        std::vector<double> threshold = matlab::linspace(threshold_from, threshold_to, count_thresholds);

        monte_carlo::engine<t_statistic_type, aftermath::probability::dist_normal> mc_engine(0.0, shift_in_mean, shift_in_mean, seed_override);

        std::vector<double> expected_time(threshold.size());
        std::vector<double> expected_aux(threshold.size());
        mc_engine.execute(n, proc, count_affected, threshold, expected_time, expected_aux, do_stop_on_reset, is_quiet, do_dump); // ARL

        matlab::toc();
    }
}
